package middlewares

import (
	"github.com/gin-gonic/gin"
	"strings"
	"sync"
	"time"
)

type TokenBucket struct {
	Capacity  int64      // 桶的容量
	Rate      float64    // 令牌放入速率
	Tokens    float64    // 当前令牌数量
	LastToken time.Time  // 上一次放令牌的时间
	Mtx       sync.Mutex // 互斥锁
}

func (tb *TokenBucket) Allow() bool {
	tb.Mtx.Lock()
	defer tb.Mtx.Unlock()
	now := time.Now()
	// 计算需要放的令牌数量
	tb.Tokens = tb.Tokens + tb.Rate*now.Sub(tb.LastToken).Seconds()
	if tb.Tokens > float64(tb.Capacity) {
		tb.Tokens = float64(tb.Capacity)
	}
	// 判断是否允许请求
	if tb.Tokens >= 1 {
		tb.Tokens--
		tb.LastToken = now
		return true
	} else {
		return false
	}
}

func LimitHandler(maxConn int64) gin.HandlerFunc {
	tb := &TokenBucket{
		Capacity:  maxConn,
		Rate:      1.0,
		Tokens:    0,
		LastToken: time.Now(),
	}
	return func(c *gin.Context) {
		if strings.Contains(c.Request.URL.Path, "/reg/pay") || strings.Contains(c.Request.URL.Path, "/reg/pay-record") || strings.Contains(c.Request.URL.Path, "/payment/test") {
			//if strings.Contains(c.Request.URL.Path, "/payment/test") {
			if !tb.Allow() {
				Logger.Println("网络繁忙，请稍后再试")
				c.AbortWithStatusJSON(503, map[string]interface{}{"msg": "网络繁忙，请稍后再试", "is_success": false})
				c.Abort()
				return
			}
			//Logger.Println("成功请求")
		}
		c.Next()
	}
}
